import json
from datetime import datetime
from typing import Dict, List, Optional, Any, Union
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class WaveDataTemplate:
    """
    脑电波数据JSON模板类
    用于定义和验证脑电波数据的JSON结构
    """
    
    @staticmethod
    def get_raw_wave_template() -> Dict[str, Any]:
        """
        获取原始脑电数据模板
        
        Returns:
            Dict: 原始脑电数据模板
        """
        return {
            "total_seconds": 0,
            "frame_rate": 128,
            "wave": {
                "frames": []
            }
        }
    
    @staticmethod
    def get_frequency_wave_template() -> Dict[str, Any]:
        """
        获取频段数据模板
        
        Returns:
            Dict: 频段数据模板
        """
        return {
            "total_seconds": 0,
            "frame_rate": 1,
            "wave": {
                "frames": []
            }
        }
    
    @staticmethod
    def get_metadata_template() -> Dict[str, Any]:
        """
        获取元数据模板
        
        Returns:
            Dict: 元数据模板
        """
        return {
            "device_info": {
                "device_name": "",
                "device_type": "EEG",
                "sampling_rate": 512,
                "transmission_rate": 128
            },
            "signal_quality": "unknown",
            "session_notes": "",
            "created_at": datetime.now().isoformat()
        }


class WaveDataManager:
    """
    脑电波数据管理器
    用于处理原始数据和频段数据的存储、解析和验证
    """
    
    def __init__(self):
        """初始化数据管理器"""
        self.raw_data_buffer = []
        self.delta_data_buffer = []
        self.theta_data_buffer = []
        self.alpha_data_buffer = []
        self.beta_data_buffer = []
        self.gamma_data_buffer = []
        self.start_time = None
        self.end_time = None
    
    def start_session(self):
        """开始数据收集会话"""
        self.start_time = datetime.now()
        self.raw_data_buffer = []
        self.delta_data_buffer = []
        self.theta_data_buffer = []
        self.alpha_data_buffer = []
        self.beta_data_buffer = []
        self.gamma_data_buffer = []
        logger.info("数据收集会话开始")
    
    def add_raw_data(self, raw_value: int, frame_index: Optional[int] = None):
        """
        添加原始脑电数据
        
        Args:
            raw_value: 原始脑电值
            frame_index: 帧索引，如果不提供则自动计算
        """
        if frame_index is None:
            frame_index = len(self.raw_data_buffer) + 1
        
        self.raw_data_buffer.append([frame_index, raw_value])
    
    def add_frequency_data(self, frequency_type: str, value: int, frame_index: Optional[int] = None):
        """
        添加频段数据
        
        Args:
            frequency_type: 频段类型 ('delta', 'theta', 'alpha', 'beta', 'gamma')
            value: 频段值
            frame_index: 帧索引，如果不提供则自动计算
        """
        if frame_index is None:
            frame_index = len(getattr(self, f"{frequency_type}_data_buffer")) + 1
        
        buffer_name = f"{frequency_type}_data_buffer"
        if hasattr(self, buffer_name):
            getattr(self, buffer_name).append([frame_index, value])
        else:
            logger.warning(f"不支持的频段类型: {frequency_type}")
    
    def end_session(self):
        """结束数据收集会话"""
        self.end_time = datetime.now()
        logger.info("数据收集会话结束")
    
    def get_session_duration(self) -> int:
        """
        获取会话持续时间（秒）
        
        Returns:
            int: 会话持续时间
        """
        if self.start_time and self.end_time:
            return int((self.end_time - self.start_time).total_seconds())
        return 0
    
    def create_raw_wave_data(self) -> Dict[str, Any]:
        """
        创建原始脑电数据JSON
        
        Returns:
            Dict: 原始脑电数据JSON
        """
        template = WaveDataTemplate.get_raw_wave_template()
        template["total_seconds"] = self.get_session_duration()
        template["wave"]["frames"] = self.raw_data_buffer.copy()
        return template
    
    def create_frequency_wave_data(self, frequency_type: str) -> Dict[str, Any]:
        """
        创建频段数据JSON
        
        Args:
            frequency_type: 频段类型
            
        Returns:
            Dict: 频段数据JSON
        """
        template = WaveDataTemplate.get_frequency_wave_template()
        template["total_seconds"] = self.get_session_duration()
        
        buffer_name = f"{frequency_type}_data_buffer"
        if hasattr(self, buffer_name):
            template["wave"]["frames"] = getattr(self, buffer_name).copy()
        else:
            template["wave"]["frames"] = []
        
        return template
    
    def create_metadata(self, device_info: Optional[Dict[str, Any]] = None, 
                       signal_quality: str = "good", 
                       session_notes: str = "") -> Dict[str, Any]:
        """
        创建元数据JSON
        
        Args:
            device_info: 设备信息
            signal_quality: 信号质量
            session_notes: 会话备注
            
        Returns:
            Dict: 元数据JSON
        """
        template = WaveDataTemplate.get_metadata_template()
        
        if device_info:
            template["device_info"].update(device_info)
        
        template["signal_quality"] = signal_quality
        template["session_notes"] = session_notes
        template["created_at"] = self.start_time.isoformat() if self.start_time else datetime.now().isoformat()
        
        return template
    
    def create_complete_session_data(self, session_id: str, 
                                   device_info: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """
        创建完整的会话数据
        
        Args:
            session_id: 会话ID
            device_info: 设备信息
            
        Returns:
            Dict: 完整的会话数据
        """
        duration = self.get_session_duration()
        
        return {
            "id": session_id,
            "session_id": session_id,
            "total_seconds": duration,
            "wave": self.create_raw_wave_data(),
            "delta_wave": self.create_frequency_wave_data("delta"),
            "theta_wave": self.create_frequency_wave_data("theta"),
            "alpha_wave": self.create_frequency_wave_data("alpha"),
            "beta_wave": self.create_frequency_wave_data("beta"),
            "gamma_wave": self.create_frequency_wave_data("gamma"),
            "metadata": self.create_metadata(device_info)
        }
    
    @staticmethod
    def parse_raw_wave_data(json_data: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
        """
        解析原始脑电数据JSON
        
        Args:
            json_data: JSON数据（字符串或字典）
            
        Returns:
            Dict: 解析后的数据
        """
        if isinstance(json_data, str):
            json_data = json.loads(json_data)
        
        template = WaveDataTemplate.get_raw_wave_template()
        template.update(json_data)
        return template
    
    @staticmethod
    def parse_frequency_wave_data(json_data: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
        """
        解析频段数据JSON
        
        Args:
            json_data: JSON数据（字符串或字典）
            
        Returns:
            Dict: 解析后的数据
        """
        if isinstance(json_data, str):
            json_data = json.loads(json_data)
        
        template = WaveDataTemplate.get_frequency_wave_template()
        template.update(json_data)
        return template
    
    @staticmethod
    def parse_metadata(json_data: Union[str, Dict[str, Any]]) -> Dict[str, Any]:
        """
        解析元数据JSON
        
        Args:
            json_data: JSON数据（字符串或字典）
            
        Returns:
            Dict: 解析后的数据
        """
        if isinstance(json_data, str):
            json_data = json.loads(json_data)
        
        template = WaveDataTemplate.get_metadata_template()
        template.update(json_data)
        return template
    
    def get_data_statistics(self) -> Dict[str, Any]:
        """
        获取数据统计信息
        
        Returns:
            Dict: 统计信息
        """
        duration = self.get_session_duration()
        
        return {
            "session_duration_seconds": duration,
            "session_duration_hours": round(duration / 3600, 2),
            "raw_data_points": len(self.raw_data_buffer),
            "delta_data_points": len(self.delta_data_buffer),
            "theta_data_points": len(self.theta_data_buffer),
            "alpha_data_points": len(self.alpha_data_buffer),
            "beta_data_points": len(self.beta_data_buffer),
            "gamma_data_points": len(self.gamma_data_buffer),
            "start_time": self.start_time.isoformat() if self.start_time else None,
            "end_time": self.end_time.isoformat() if self.end_time else None
        }
    
    def validate_data(self) -> Dict[str, bool]:
        """
        验证数据完整性
        
        Returns:
            Dict: 验证结果
        """
        duration = self.get_session_duration()
        expected_raw_frames = duration * 128  # 128Hz采样率
        expected_freq_frames = duration * 1   # 1Hz频段数据
        
        validation_results = {
            "has_start_time": self.start_time is not None,
            "has_end_time": self.end_time is not None,
            "duration_valid": duration > 0,
            "raw_data_complete": len(self.raw_data_buffer) >= expected_raw_frames * 0.9,  # 允许10%误差
            "delta_data_complete": len(self.delta_data_buffer) >= expected_freq_frames * 0.9,
            "theta_data_complete": len(self.theta_data_buffer) >= expected_freq_frames * 0.9,
            "alpha_data_complete": len(self.alpha_data_buffer) >= expected_freq_frames * 0.9,
            "beta_data_complete": len(self.beta_data_buffer) >= expected_freq_frames * 0.9,
            "gamma_data_complete": len(self.gamma_data_buffer) >= expected_freq_frames * 0.9
        }
        
        return validation_results


# 便捷函数
def create_wave_data_from_raw(raw_data: List[List[int]], 
                             duration_seconds: int,
                             frame_rate: int = 128) -> Dict[str, Any]:
    """
    从原始数据创建脑电数据JSON
    
    Args:
        raw_data: 原始数据列表 [[frame_index, value], ...]
        duration_seconds: 持续时间（秒）
        frame_rate: 帧率
        
    Returns:
        Dict: 脑电数据JSON
    """
    template = WaveDataTemplate.get_raw_wave_template()
    template["total_seconds"] = duration_seconds
    template["frame_rate"] = frame_rate
    template["wave"]["frames"] = raw_data
    return template


def create_frequency_data_from_values(values: List[List[int]], 
                                     duration_seconds: int,
                                     frame_rate: int = 1) -> Dict[str, Any]:
    """
    从频段值创建频段数据JSON
    
    Args:
        values: 频段值列表 [[frame_index, value], ...]
        duration_seconds: 持续时间（秒）
        frame_rate: 帧率
        
    Returns:
        Dict: 频段数据JSON
    """
    template = WaveDataTemplate.get_frequency_wave_template()
    template["total_seconds"] = duration_seconds
    template["frame_rate"] = frame_rate
    template["wave"]["frames"] = values
    return template


# 使用示例
if __name__ == "__main__":
    # 创建数据管理器
    manager = WaveDataManager()
    
    # 开始会话
    manager.start_session()
    
    # 模拟添加数据
    for i in range(1000):
        manager.add_raw_data(raw_value=1000 + i, frame_index=i+1)
        
        if i % 128 == 0:  # 每128个原始数据点添加一次频段数据
            freq_index = i // 128 + 1
            manager.add_frequency_data("delta", 200 + freq_index, freq_index)
            manager.add_frequency_data("theta", 300 + freq_index, freq_index)
            manager.add_frequency_data("alpha", 400 + freq_index, freq_index)
            manager.add_frequency_data("beta", 500 + freq_index, freq_index)
            manager.add_frequency_data("gamma", 600 + freq_index, freq_index)
    
    # 结束会话
    manager.end_session()
    
    # 创建完整数据
    session_data = manager.create_complete_session_data(
        session_id="test_session_001",
        device_info={
            "device_name": "NeuroSky MindWave",
            "device_type": "EEG"
        }
    )
    
    # 输出统计信息
    stats = manager.get_data_statistics()

    
    # 验证数据
    validation = manager.validate_data()

    
    # 测试解析功能
    raw_wave_json = json.dumps(session_data["wave"])
    parsed_raw = WaveDataManager.parse_raw_wave_data(raw_wave_json)